import java.util.Arrays;

public class Interview40 {
    /**
     * 输入整数数组 arr ，找出其中最小的 k 个数。
     * 例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。
     * <p>
     * 示例 1：
     * <p>
     * 输入：arr = [3,2,1], k = 2
     * 输出：[1,2] 或者 [2,1]
     */
    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr == null || arr.length == 0 || k <=0) return new int[0];
        int start = 0, end = arr.length - 1;
        int index = partition(arr, start, end);
        while (index != k - 1) {
            if (index < k - 1) {
                start = index + 1;
            } else {
                end = index - 1;
            }
            index = partition(arr, start, end);
        }
        return Arrays.copyOfRange(arr, 0, index + 1);
    }

    public int partition(int[] arr, int lo, int hi) {
        int index = lo + (int) (Math.random() * (hi - lo + 1));
        swap(arr, index, hi);
        int small = lo - 1;
        while (lo <= hi) {
            if (arr[lo] < arr[hi]) {
                small++;
                swap(arr, small, lo);
            }
            lo++;
        }
        small++;
        swap(arr, small, hi);
        return small;
    }

    public void swap(int[] arr, int i, int j) {
        int temp = arr[j];
        arr[j] = arr[i];
        arr[i] = temp;
    }
}
